#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-18


import threading
import time
import cv2
import os
import json
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg, BaseNode, get_extra_args
import argparse
from datetime import datetime


class SpireViewSaverNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.output_spire_dir = self.get_param("output_spire_dir", "G:/dataset/240918_false_detections")
        self.imgres_topic = self.get_param("imgres_topic", "")
        self.save_empty = self.get_param("save_empty", 0)
        self.params_help()

        input_topic = '/' + job_name + '/detector/image_results'
        if len(self.imgres_topic) > 0:
            input_topic = self.imgres_topic

        self._image_reader = Subscriber(
            input_topic, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )

        self.image_queue = Queue()
        self.result_queue = Queue()
        self.queue_pool.append(self.image_queue)
        self.queue_pool.append(self.result_queue)
        self.run_time = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
        self.frame_cnt = 0
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()

    def image_callback(self, msg):
        img = sms2cvimg(msg)
        self.image_queue.put({'msg': msg, 'img': img})

    def run(self):
        while self.is_running():
            msg_img = self.image_queue.get(block=True)
            if msg_img is None:
                break
            img, msg1 = msg_img['img'], msg_img['msg']

            if 'spirecv_msgs::2DTargets' in msg1:
                self.frame_cnt += 1
                msg2 = msg1['spirecv_msgs::2DTargets']
                image_path = os.path.join(self.output_spire_dir, 'scaled_images')
                anno_path = os.path.join(self.output_spire_dir, 'annotations')
                if not os.path.exists(image_path):
                    os.makedirs(image_path)
                if not os.path.exists(anno_path):
                    os.makedirs(anno_path)

                spire_anno = {}
                spire_anno['file_name'] = msg2['file_name'] if len(msg2['file_name']) > 0 else "{}_{}.jpg".format(self.run_time, str(self.frame_cnt).zfill(7))
                spire_anno['height'] = img.shape[0]
                spire_anno['width'] = img.shape[1]
                spire_anno['annos'] = []
                for target in msg2['targets']:
                    anno = {}
                    anno['bbox'] = target['bbox']
                    anno['score'] = target['score']
                    anno['category_name'] = target['category_name']
                    anno['area'] = anno['bbox'][2] * anno['bbox'][3]
                    spire_anno['annos'].append(anno)

                if len(spire_anno['annos']) or self.save_empty > 0:
                    output_fn = os.path.join(anno_path, spire_anno['file_name'] + '.json')
                    with open(output_fn, "w") as f:
                        json.dump(spire_anno, f)

                    cv2.imwrite(os.path.join(image_path, spire_anno['file_name']), img)
                    print("Saved: {}".format(spire_anno['file_name']))

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("config:", args.config)
    print("job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = SpireViewSaverNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()
